# -*- coding: utf-8 -*-
"""
Created on Wed Apr 14 16:18:15 2021
Description : 深度信念网络
@author: Lenovo-pc
"""

import warnings
warnings.filterwarnings("ignore")

# import sys
# import os
import pickle

import numpy as np
from deep_learning_algs.rbm_binary import rbm_binary
from deep_learning_algs.dnn import dnn

class dbn_binary:
    def __init__(self):
        pass
        
    def weight(self,datasets,hidden_units = [500,100] ):

        h = datasets
        
        self.num_visible = datasets.shape[1]
        self.hidden_units = hidden_units # 每个隐层的单元数
        self.num_layers = len(hidden_units)
        layers = []
        i = self.num_visible
        for j in hidden_units:
            layers.append(rbm_binary(i,j))
            i = j
        self.layers = layers #受限玻尔兹曼机 训练结果
        #self.weights = np.array([weights])
        self.weights = []
        
        for n, layer in enumerate(self.layers):
            
            layer.train(h,batch_size=6,max_epoch=50)
            self.weights.append(layer.weights)
            if n+1<self.num_layers:
                h = layer.get_hidden(h)
            else:
                self.last_h = layer.get_hidden(h)

        # 评估
        sample = self.get_visble(self.last_h)
        #print(sample)
        error = np.sum((datasets-sample)**2)
        print("DBN : error is %s" % (error))
    
    def save_weight(self,save_path):
        weight = {"num_layers":self.num_layers,
                  "layers":self.layers,
                  "weights":self.weights
            }
        # dump功能
        # dump 将数据通过特殊的形式转换为只有python语言认识的字符串，并写入文件
        with open(save_path, 'w') as f:
            pickle.dump(weight, f)
    
    def load_weight(self,save_path):
        with open(save_path, 'r') as f:
            weight = pickle.load(f)
        self.layers = weight["layers"]
        self.weights = weight["weights"]
        self.num_layers = weight["num_layers"]
        
    def fine_tuning(self,train_data,train_labels):
        self.dnn_model = dnn()
        self.dnn_model.train(train_data,train_labels)
      
    
    def train(self,train_data,train_labels):
        assert (len(train_data)==len(train_labels)) and\
            len(train_data)>0 and len(train_labels)>0,\
        'dbn_binary ERROR:train input'
        self.weight(train_data)
        self.fine_tuning(self.last_h,train_labels)
        
        
    def predict_data(self,input_data):
        h_layer = self.get_last_hidden(input_data)
        pred_data = self.dnn_model.predict_data(h_layer)
        pred_data = pred_data.reshape(len(input_data))
        return pred_data
    
    def estimate_cor(self,test_labels,predict_test):
        cor = np.corrcoef(test_labels,predict_test)
        return cor
        

    def get_last_hidden(self,data):
        visible = data
        for layer in self.layers:
            visible = layer.get_hidden(visible)

        hidden = visible
        return hidden

    def get_visble(self,data):
        hidden = data
        for layer in reversed(self.layers):
            hidden = layer.get_visible(hidden)

        visible = hidden
        return visible


if __name__ == '__main__':
    r = dbn_binary()
    training_data = np.array([[1,1,1,0,0,0],
                              [1,0,1,0,0,0],
                              [1,1,1,0,0,0],
                              [0,0,1,1,1,0],
                              [0,0,1,1,0,0],
                              [0,0,1,1,1,0]])
    r.pre_train(training_data)
    #print(r.weights)

    user = np.array([[0,0,0,1,1,0]])
    print(r.get_last_hidden(user))
    sample = r.get_visble(r.get_last_hidden(user))
    print(sample)